# JWT工具类
import jwt
import datetime
import os
from functools import wraps
from flask import current_app, request, jsonify
from typing import Dict, Any, Optional, Union

class JWTUtils:
    """JWT工具类，用于处理JWT令牌的生成、验证和解析"""
    
    @staticmethod
    def get_secret_key():
        """
        从配置或环境变量中获取JWT密钥
        
        Returns:
            JWT密钥字符串
        """
        # 首先尝试从Flask应用配置中获取
        secret_key = current_app.config.get('JWT_SECRET_KEY') or current_app.config.get('SECRET_KEY')
        
        # 如果配置中没有，尝试从环境变量中获取
        if secret_key is None:
            secret_key = os.environ.get('JWT_SECRET_KEY') or os.environ.get('SECRET_KEY')
        
        # 如果仍然没有密钥，抛出异常
        if secret_key is None:
            raise ValueError("JWT密钥未配置，请设置JWT_SECRET_KEY环境变量或应用配置")
        
        # 确保密钥是字符串类型
        if not isinstance(secret_key, str):
            secret_key = str(secret_key)
            
        return secret_key
    
    @staticmethod
    def generate_token(payload: Dict[str, Any], expires_in: int = None) -> str:
        """
        生成JWT令牌
        
        Args:
            payload: 要编码的数据载荷
            expires_in: 过期时间（秒），默认为1小时
            
        Returns:
            JWT令牌字符串
        """
        if expires_in is None:
            expires_in = current_app.config.get('JWT_EXPIRATION_HOURS', 1) * 3600
            
        # 设置过期时间
        exp = datetime.datetime.utcnow() + datetime.timedelta(seconds=expires_in)
        
        # 添加标准声明
        payload.update({
            'exp': exp,
            'iat': datetime.datetime.utcnow(),  # 签发时间
            'iss': current_app.config.get('JWT_ISSUER', 'FlaskApp')  # 签发者
        })
        
        # 生成JWT
        secret_key = JWTUtils.get_secret_key()
        algorithm = current_app.config.get('JWT_ALGORITHM', 'HS256')
        
        token = jwt.encode(payload, secret_key, algorithm=algorithm)
        
        # 确保返回字符串类型（PyJWT 2.0+返回字符串，旧版本返回字节）
        if isinstance(token, bytes):
            return token.decode('utf-8')
        return token
    
    @staticmethod
    def verify_token(token: str) -> Optional[Dict[str, Any]]:
        """
        验证JWT令牌
        
        Args:
            token: JWT令牌字符串
            
        Returns:
            解码后的载荷，验证失败返回None
        """
        try:
            secret_key = JWTUtils.get_secret_key()
            algorithm = current_app.config.get('JWT_ALGORITHM', 'HS256')
            
            # 验证并解码JWT
            payload = jwt.decode(
                token, 
                secret_key, 
                algorithms=[algorithm],
                options={
                    'verify_exp': True,  # 验证过期时间
                    'verify_iat': True,  # 验证签发时间
                    'verify_iss': True   # 验证签发者
                },
                issuer=current_app.config.get('JWT_ISSUER', 'FlaskApp')
            )
            
            return payload
        except jwt.ExpiredSignatureError:
            # 令牌已过期
            return None
        except jwt.InvalidTokenError:
            # 令牌无效
            return None
    
    @staticmethod
    def refresh_token(token: str) -> Optional[str]:
        """
        刷新JWT令牌
        
        Args:
            token: 原始JWT令牌
            
        Returns:
            新的JWT令牌，刷新失败返回None
        """
        try:
            # 验证原始令牌（不检查过期时间）
            secret_key = JWTUtils.get_secret_key()
            algorithm = current_app.config.get('JWT_ALGORITHM', 'HS256')
            
            payload = jwt.decode(
                token, 
                secret_key, 
                algorithms=[algorithm],
                options={'verify_exp': False}  # 不验证过期时间
            )
            
            # 移除时间相关声明
            payload.pop('exp', None)
            payload.pop('iat', None)
            payload.pop('iss', None)
            
            # 生成新令牌
            return JWTUtils.generate_token(payload)
        except jwt.InvalidTokenError:
            # 令牌无效
            return None
    
    @staticmethod
    def get_token_from_request() -> Optional[str]:
        """
        从请求中提取JWT令牌
        
        Returns:
            JWT令牌字符串，未找到返回None
        """
        auth_header = request.headers.get('Authorization')
        if auth_header:
            # 通常格式为 "Bearer <token>"
            parts = auth_header.split()
            if len(parts) == 2 and parts[0].lower() == 'bearer':
                return parts[1]
        
        # 也可以从查询参数中获取
        return request.args.get('token')
    
    @staticmethod
    def is_token_expired(token: str) -> bool:
        """
        检查令牌是否过期
        
        Args:
            token: JWT令牌字符串
            
        Returns:
            True表示已过期，False表示未过期或无效
        """
        try:
            secret_key = JWTUtils.get_secret_key()
            algorithm = current_app.config.get('JWT_ALGORITHM', 'HS256')
            
            # 只验证过期时间
            jwt.decode(
                token, 
                secret_key, 
                algorithms=[algorithm],
                options={'verify_exp': True}
            )
            return False
        except jwt.ExpiredSignatureError:
            return True
        except jwt.InvalidTokenError:
            return True
    
    @staticmethod
    def decode_token_without_verification(token: str) -> Optional[Dict[str, Any]]:
        """
        不验证签名的情况下解码JWT令牌（仅用于调试）
        
        Args:
            token: JWT令牌字符串
            
        Returns:
            解码后的载荷，解码失败返回None
        """
        try:
            # 不验证签名，只解码
            return jwt.decode(token, options={'verify_signature': False})
        except:
            return None


def token_required(f):
    """
    JWT认证装饰器，用于保护需要认证的路由
    
    使用方法:
    @token_required
    def protected_route():
        return jsonify({'message': 'This is a protected route'})
    """
    @wraps(f)
    def decorated(*args, **kwargs):
        # 从请求中获取令牌
        token = JWTUtils.get_token_from_request()
        
        if not token:
            return jsonify({
                'code': 401,
                'msg': '令牌缺失',
                'data': None
            }), 401
        
        # 验证令牌
        payload = JWTUtils.verify_token(token)
        if not payload:
            return jsonify({
                'code': 401,
                'msg': '令牌无效或已过期',
                'data': None
            }), 401
        
        # 将用户信息添加到请求上下文
        request.current_user = payload
        
        return f(*args, **kwargs)
    
    return decorated


def admin_required(f):
    """
    管理员权限装饰器，用于保护需要管理员权限的路由
    
    使用方法:
    @admin_required
    def admin_route():
        return jsonify({'message': 'This is an admin route'})
    """
    @wraps(f)
    def decorated(*args, **kwargs):
        # 先进行JWT认证
        token = JWTUtils.get_token_from_request()
        
        if not token:
            return jsonify({
                'code': 401,
                'msg': '令牌缺失',
                'data': None
            }), 401
        
        payload = JWTUtils.verify_token(token)
        if not payload:
            return jsonify({
                'code': 401,
                'msg': '令牌无效或已过期',
                'data': None
            }), 401
        
        # 检查是否为管理员
        if not payload.get('is_admin', False):
            return jsonify({
                'code': 403,
                'msg': '需要管理员权限',
                'data': None
            }), 403
        
        # 将用户信息添加到请求上下文
        request.current_user = payload
        
        return f(*args, **kwargs)
    
    return decorated